#%%
import numpy as np, matplotlib.pyplot as plt, seaborn as sns, numpy as np, sys, os, tqdm, lmfit as lm
import pandas as pd, scipy.stats as stats, scipy.optimize as op, scipy.interpolate as ip, pathlib
import grain_shape_effects_data_proc_adjusted_stress as grain_data
import matplotlib.ticker as mticker
from matplotlib.ticker import FormatStrFormatter

%matplotlib inline
plt.style.use('seaborn-whitegrid')
sns.set_context('paper')
sns.set_style('ticks')
pd.options.mode.chained_assignment = None
label = {'gb': 'Spheres', 'ng': 'Natural gravel', 'gg': 'Rounded chips', 'ls': 'Rectangular prisms', 'oc': 'Faceted ellipsoids'}
color = {'gb': sns.xkcd_rgb['cobalt blue'], 'gg': sns.xkcd_rgb['cerulean'], 'oc': sns.xkcd_rgb['blue green'], 'ng': sns.xkcd_rgb['maize'], 'ls': sns.xkcd_rgb['blood orange'], 'bm': 'y'}
marker = {'gb': 'o', 'gg': 'D', 'oc': '*', 'ng': 'v', 'ls': 's'}
size = {'gb': 9, 'gg': 9, 'oc': 18, 'ng': 12, 'ls': 9}

main = grain_data.Fits(grains=['gb', 'oc', 'gg', 'ng', 'ls'])
figname = 'print'

#%%
fig2 = plt.figure(constrained_layout=True, figsize=(20,12))
gs = fig2.add_gridspec(4,3)

f2_ax1 = fig2.add_subplot(gs[:2, 1])
main.plot_standard(ax=f2_ax1)
f2_ax1.xaxis.set_major_formatter(mticker.ScalarFormatter())
f2_ax1.xaxis.set_minor_formatter(mticker.ScalarFormatter())
f2_ax1.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
f2_ax1.tick_params(axis='x', which='major', labelsize=15, rotation = 65)
f2_ax1.tick_params(axis='x', which='minor', labelsize=15, rotation = 65)
f2_ax1.text(0.045, .8, 'a', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax2 = fig2.add_subplot(gs[:2, 0])
main.plot_CDx(ax=f2_ax2)

f2_ax2.xaxis.set_major_formatter(mticker.ScalarFormatter())
f2_ax2.xaxis.set_minor_formatter(mticker.ScalarFormatter())
f2_ax2.xaxis.set_major_formatter(FormatStrFormatter('%.2f'))
f2_ax2.tick_params(axis='x', which='major', labelsize=15, rotation = 65)
f2_ax2.tick_params(axis='x', which='minor', labelsize=15, rotation = 65)
f2_ax2.text(0.025, .8, 'b', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax3 = fig2.add_subplot(gs[2, 0])
main.plot_coeff(ax=f2_ax3)
f2_ax3.text(2.8, 12, 'd', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax4 = fig2.add_subplot(gs[3, 0], sharex=f2_ax3)
main.plot_taucs(ax=f2_ax4)
f2_ax4.text(2.8, 0.12, 'e', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax5 = fig2.add_subplot(gs[:2, 2])
main.plot_mus(ax=f2_ax5)
f2_ax5.xaxis.set_major_formatter(mticker.ScalarFormatter())
f2_ax5.xaxis.set_minor_formatter(mticker.ScalarFormatter())
f2_ax5.xaxis.set_major_formatter(FormatStrFormatter('%.2f'))
f2_ax5.tick_params(axis='x', which='major', labelsize=15, rotation = 65)
f2_ax5.tick_params(axis='x', which='minor', labelsize=15, rotation = 65)
f2_ax5.text(0.025, .8, 'c', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax6 = fig2.add_subplot(gs[2, 2])
main.plot_coeff_2(ax=f2_ax6)
f2_ax6.text(2.8, 40, 'f', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax7 = fig2.add_subplot(gs[3, 2], sharex=f2_ax6)
main.plot_taucs_2(ax=f2_ax7)
f2_ax7.text(2.8, 0.05, 'g', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

f2_ax8 = fig2.add_subplot(gs[2:, 1])
main.plot_final_new(ax=f2_ax8)
f2_ax8.xaxis.set_major_formatter(mticker.ScalarFormatter())
f2_ax8.xaxis.set_minor_formatter(mticker.ScalarFormatter())
f2_ax8.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
f2_ax8.tick_params(axis='x', which='major', labelsize=15, rotation = 65)
f2_ax8.tick_params(axis='x', which='minor', labelsize=15, rotation = 65)
f2_ax8.text(0.045, .8, 'h', fontsize=24, bbox=dict(facecolor='white', edgecolor='none', pad=4))

sns.despine()
plt.tight_layout()
plt.savefig('fig_2b%s.png'%figname, dpi=300)

# %%
